


"""
去边留芯分块块 函数

将输入张量通过 “去边留芯” 的方式，调节到可以被patch_size刚好裁剪完的形式，裁剪成一堆小块块
举例：
    输入尺寸：32个单通道[121, 145, 121]的长方形块块：[32, 1, 121, 145, 121]
    patch_size:30
    输出：每个大块块，拿到最中心的[120, 120, 120]的块块后裁剪成 4x4x4 个 [30,30,30]的小块块，再把每一组（共32组）小块块放到list里
    得到[32, 4x4x4, 30, 30, 30]
"""
import torch
from einops import rearrange
from torch import nn


class Patches_CropSelected(nn.Module):  # for brain neuroimaging (edge patch will be excluded)
    def __init__(self, input_shape, patch_size, crop_mode="center"):
        super().__init__()

        assert (len(input_shape) == 3)
        assert (isinstance(patch_size, int) or (isinstance(patch_size, list) and len(patch_size) == 3))

        if isinstance(patch_size, int):
            self.patch_size = [patch_size, patch_size, patch_size]
        else:
            self.patch_size = patch_size

        self.new_ipt_shape = [input_shape[idx] % self.patch_size[idx] for idx in range(3)]
        self.edge1_shape = [self.new_ipt_shape[idx] // 2 for idx in range(3)]
        self.edge2_shape = [self.new_ipt_shape[idx] - self.edge1_shape[idx] for idx in range(3)]

        for i in range(3):
            if self.edge2_shape[i] == 0:
                self.edge2_shape[i] = None
            else:
                self.edge2_shape[i] = -self.edge2_shape[i]
        self.patches_shape = torch.tensor(input_shape) // self.patch_size[0]

    def forward(self, x):
        B = x.shape[0]

        x = x[:, :,
            self.edge1_shape[0]:self.edge2_shape[0],
            self.edge1_shape[1]:self.edge2_shape[1],
            self.edge1_shape[2]:self.edge2_shape[2]]

        patchs_output = rearrange(x, 'b c (h1 ph) (w1 pw) (d1 pd) -> (b h1 w1 d1) c ph pw pd',
                                  ph=self.patch_size[0], pw=self.patch_size[1], pd=self.patch_size[2])

        patchs_output = rearrange(patchs_output, '(b p) ... -> b p ...', b=B)

        patchs_output = rearrange(patchs_output, 'b p c ... -> b (p c) ...', b=B)

        return patchs_output


if __name__ == '__main__':
    model = Patches_CropSelected([50, 125, 100], 25)
    x = torch.randn((3, 1, 50, 125, 100))
    y = model(x)
    print(y.shape)
    # torch.Size([3, 40, 25, 25, 25])
    print(model.patches_shape)
    # tensor([2, 5, 4])
